use std::fs::{create_dir_all, File};
use std::io::Write;
use std::path::{Path, PathBuf};
use log::{error, info};
use crate::utils::{copy_file, get_image_size, get_max_num};
use crate::converter::bbox2yolo;
use crate::parser::parse_bbox_from_filename;
use crate::YoloData;

fn write_yolo_data(yolo_data: YoloData, dst_filename: &str, class_id: i32) -> Result<(), Box<dyn std::error::Error>>{
    let mut file = File::create(&dst_filename)?;
    file.write_all(format!("{} {} {} {} {}", class_id, yolo_data.x_center, yolo_data.y_center, yolo_data.width, yolo_data.height).as_bytes())?;
    Ok(())
}


pub fn create_yolo_dataset(image_files: Vec<(String,String)>, class_id: i32, train_path_images: PathBuf, train_path_labels: PathBuf) -> Result<(), Box<dyn std::error::Error>>{
    create_dir_all(&train_path_images)?;
    create_dir_all(&train_path_labels)?;
    // 创建训练集的文件编号，默认从1开始，如果文件已经存在，则从文件编号最大值开始
    let mut i = get_max_num(&train_path_images);
    for image_file in image_files {
        // 创建新的文件名, 格式为 [类别编号]_[文件编号].jpg
        let new_filename = format!("{}_{:06}", class_id, i);

        let bbox = match parse_bbox_from_filename(&image_file.1) {
            Ok(bbox) => bbox,
            Err(e) => {
                error!("{}: Failed to parse bbox from filename '{}' (path: {}): {}",
                       new_filename, image_file.1, image_file.0, e);
                continue;
            }
        };
        // 获取图像尺寸
        let image_size = match get_image_size(&image_file.0) {
            Ok(size) => size,
            Err(e) => {
                error!("{}: Failed to get image size for file '{}' (path: {}): {}",
                       new_filename, image_file.0, image_file.0, e);
                continue;
            }
        };
        // 解析YOLO数据
        let yolo_data = match bbox2yolo(bbox, image_size, &image_file.0) {
            Ok(data) => data,
            Err(e) => {
                error!("{}: Failed to parse YOLO data : {}", new_filename, e);
                continue;
            }
        };

        // 获取原文件的扩展名
        let src_path = Path::new(&image_file.0);
        let extension = src_path.extension().and_then(|s| s.to_str()).unwrap_or("jpg");

        // 正确构建目标文件路径
        let dst_image_path = train_path_images.join(format!("{}.{}", new_filename, extension));
        let dst_label_path = train_path_labels.join(format!("{}.txt", new_filename));

        // 复制图像文件
        if let Err(e) = copy_file(&image_file.0, dst_image_path.to_str().unwrap()) {
            error!("{}: Failed to copy image file '{}' to '{}': {}",
                   new_filename, image_file.0, dst_image_path.display(), e);
            continue;
        }
        // 保存标签文件
        if let Err(e) = write_yolo_data(yolo_data, dst_label_path.to_str().unwrap(), class_id) {
            error!("{}: Failed to save YOLO data for file '{}' to '{}': {}",
                   new_filename, image_file.1, dst_label_path.display(), e);
            continue;
        }

        info!("Processed {} ==> {}", new_filename, image_file.1);
        i += 1;
    }
    Ok(())
}